package com.zhz.elasticjob.strategy;

import com.dangdang.ddframe.job.lite.api.strategy.JobInstance;
import com.dangdang.ddframe.job.lite.api.strategy.JobShardingStrategy;
import org.springframework.util.CollectionUtils;

import java.util.*;

/**
 * @program: schedule
 * @description: 是自定义分片策略
 * @author: zhanghz001
 * @create: 2021-08-05 21:07
 **/
public class ZhzJobShardingStrategy implements JobShardingStrategy {
    
    @Override
    public Map<JobInstance, List<Integer>> sharding(List<JobInstance> jobInstances,
                                                    String jobName,
                                                    int shardingTotalCount) {
        if (CollectionUtils.isEmpty(jobInstances)) {
            return Collections.emptyMap();
        }
        
        //服务和分片列表对应关系
        HashMap<JobInstance, List<Integer>> allMap = new HashMap<>();
        
        //所有的分片项列表 1,2,3....
        List<Integer> allShardingItem = new ArrayList<>();
        for (int i = 0; i < shardingTotalCount; i++) {
            allShardingItem.add(i);
        }
        
        //迭代服务列表
        for (int i = 0; i < jobInstances.size(); i++) {
            JobInstance jobInstance = jobInstances.get(i);
            //给每个服务指定分片项的集合
            List<Integer> shardingItem = getShardingItem(i, jobInstances.size(), allShardingItem);
            allMap.put(jobInstance, shardingItem);
        }
        //
        return allMap;
    }
    
    private List<Integer> getShardingItem(int jobInstanceIndex,
                                          int size, List<Integer> allShardingItem) {
        List<Integer> shardingItem = new ArrayList<>();
        
        for (Integer item : allShardingItem) {
            if (item % size == jobInstanceIndex) {
                shardingItem.add(item);
            }
        }
        return shardingItem;
    }
}
